from mark_agent import Position_mark
import numpy as np
import reedsolo as rs
import random
import json
import sqlite3
import re
from collections.abc import Iterable
from typing import Dict, List, Any, Union

class PrivacyLevelCalculator:
    def __init__(self, database_path: str = "privacy_rules.db"):
        self.db_path = database_path
        self._init_database()
    
    def _init_database(self):
        """初始化数据库和隐私规则表"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 创建隐私规则表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS privacy_rules (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                pattern TEXT NOT NULL,
                category TEXT NOT NULL,
                importance_level INTEGER NOT NULL,
                weight FLOAT NOT NULL,
                description TEXT,
                is_active BOOLEAN DEFAULT 1,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 插入默认规则
        default_rules = [
            # 身份证号
            (r'\b\d{17}[\dXx]\b', '身份证号', 5, 1.0, '身份证号码'),
            (r'\b\d{6}\d{8}\d{4}\b', '身份证号', 5, 1.0, '身份证号码格式'),
            
            # 手机号
            (r'\b1[3-9]\d{9}\b', '手机号', 4, 0.9, '手机号码'),
            (r'\b\d{3}-\d{4}-\d{4}\b', '手机号', 4, 0.9, '带分隔符手机号'),
            
            # 银行卡号
            (r'\b\d{16,19}\b', '银行卡', 4, 0.8, '银行卡号'),
            
            # 邮箱
            (r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b', '邮箱', 3, 0.7, '邮箱地址'),
            
            # 密码相关
            (r'password[\s]*[=:][\s]*[^\s]+', '密码', 5, 1.0, '密码字段'),
            (r'\b(pwd|passwd)[\s]*[=:][\s]*[^\s]+', '密码', 5, 1.0, '密码缩写'),
            
            # token相关
            (r'\b(eyJ[A-Za-z0-9-_=]+\.[A-Za-z0-9-_=]+\.?[A-Za-z0-9-_.+/=]*)\b', 'Token', 4, 0.9, 'JWT Token'),
            (r'\b[a-f0-9]{32}\b', 'Token', 4, 0.8, 'MD5哈希'),
            (r'\b[a-f0-9]{40}\b', 'Token', 4, 0.8, 'SHA1哈希'),
            (r'\b[a-f0-9]{64}\b', 'Token', 4, 0.9, 'SHA256哈希'),
            
            # 地址信息
            (r'.*(省|市|区|县|街道|路|号).*', '地址', 2, 0.5, '地理地址'),
            
            # 日期（生日）
            (r'\b(19|20)\d{2}[-/年](0[1-9]|1[0-2])[-/月](0[1-9]|[12][0-9]|3[01])日?\b', '生日', 3, 0.6, '生日日期'),
            
            # IP地址
            (r'\b(?:\d{1,3}\.){3}\d{1,3}\b', 'IP地址', 2, 0.4, 'IP地址'),
            
            # 金额
            (r'[¥￥$]?\s*\d+(?:\.\d{2})?\s*[元美元]?', '金额', 2, 0.3, '金融金额'),
        ]
        
        cursor.execute("SELECT COUNT(*) FROM privacy_rules")
        if cursor.fetchone()[0] == 0:
            cursor.executemany('''
                INSERT INTO privacy_rules (pattern, category, importance_level, weight, description)
                VALUES (?, ?, ?, ?, ?)
            ''', default_rules)
        
        conn.commit()
        conn.close()
    
    def get_active_rules(self) -> List[Dict]:
        """获取所有激活的隐私规则"""
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT pattern, category, importance_level, weight, description 
            FROM privacy_rules 
            WHERE is_active = 1
        ''')
        
        rules = [dict(row) for row in cursor.fetchall()]
        conn.close()
        return rules
    
    def calculate_privacy_level(self, data: Any) -> int:
        """
        计算数据的隐私等级（1-5级）
        基于正则匹配value值的重要性进行加权计算
        """
        # 将数据转换为可处理的格式
        processed_data = self._preprocess_data(data)
        
        # 获取所有激活的规则
        rules = self.get_active_rules()
        
        total_score = 0.0
        matched_patterns = []
        
        # 递归遍历数据，匹配value
        total_score = self._traverse_and_match(processed_data, rules, matched_patterns)
        
        # 根据总分计算隐私等级
        privacy_level = self._score_to_level(total_score)
        
        return privacy_level
    
    def _preprocess_data(self, data: Any) -> Any:
        """预处理数据，确保可以递归遍历"""
        if isinstance(data, (dict, list, str, int, float, bool)):
            return data
        else:
            # 尝试转换为字符串
            try:
                return str(data)
            except:
                return ""
    
    def _traverse_and_match(self, data: Any, rules: List[Dict],
                            matched_patterns: List[Dict]) -> float:
        """返回本轮累加后的分数"""
        score = 0.0
        if isinstance(data, dict):
            for v in data.values():
                score += self._traverse_and_match(v, rules, matched_patterns)
        elif isinstance(data, list):
            for item in data:
                score += self._traverse_and_match(item, rules, matched_patterns)
        elif isinstance(data, str):
        # 对字符串值进行匹配
            for rule in rules:
                if re.search(rule['pattern'], data, re.IGNORECASE):
                    s = rule['importance_level'] * rule['weight']
                    score += s
                    matched_patterns.append({
                        'pattern': rule['pattern'],
                        'category': rule['category'],
                        'score': s,
                        'matched_value': data[:50] + "..." if len(data) > 50 else data
                    })
        return score
    
    def _score_to_level(self, score: float) -> int:
        """将分数转换为隐私等级（1-5）"""
        if score == 0:
            return 1
        elif score < 3:
            return 2
        elif score < 7:
            return 3
        elif score < 12:
            return 4
        else:
            return 5

# 初始化隐私等级计算器
privacy_calculator = PrivacyLevelCalculator()

class TokenEncoder:
    def __init__(self):
        self.levels = {
            1: 4, 2: 8, 3: 12, 4: 16, 5: 20
        }
    
    def encode_token(self, token, privacy_level):
        """编码token，自动处理长度问题"""
        if privacy_level not in self.levels:
            raise ValueError("隐私等级必须在1-5之间")
        
        # 处理token长度
        if len(token) < 20:
            token = token.ljust(20, '0')
        elif len(token) > 20:
            token = token[:20]
        
        ecc_symbols = self.levels[privacy_level]
        token_bytes = token.encode('utf-8')
        codec = rs.RSCodec(ecc_symbols)
        
        try:
            encoded_data = codec.encode(token_bytes)
            return encoded_data
        except Exception as e:
            raise ValueError(f"Token编码失败: {str(e)}")
    
    def decode_token(self, encoded_data, privacy_level):
        """解码token，包含错误处理"""
        if privacy_level not in self.levels:
            raise ValueError("隐私等级必须在1-5之间")
        
        ecc_symbols = self.levels[privacy_level]
        codec = rs.RSCodec(ecc_symbols)
        
        try:
            decoded_data, _, _ = codec.decode(encoded_data)
            token = decoded_data.decode('utf-8').rstrip('0')  # 移除填充
            return token
        except Exception as e:
            raise ValueError(f"Token解码失败: {str(e)}")

token_encoder = TokenEncoder()

def calculate_privacy_level(original_data):
    """计算隐私等级 - 使用新的基于正则规则的方法"""
    return privacy_calculator.calculate_privacy_level(original_data)

def generate(original_data, token):
    """生成标识比特"""
    privacy_level = calculate_privacy_level(original_data)
    
    try:
        encoded_token = token_encoder.encode_token(token, privacy_level)
        mark_bits = ''.join(format(byte, '08b') for byte in encoded_token)
        return mark_bits, original_data, privacy_level
    except Exception as e:
        raise ValueError(f"生成标识失败: {str(e)}")

def decode(mark_bits, original_data):
    """从标识比特解码token，自动从原始数据计算隐私等级"""
    try:
        # 从原始数据计算隐私等级
        privacy_level = calculate_privacy_level(original_data)
        
        bit_length = len(mark_bits)
        if bit_length % 8 != 0:
            raise ValueError("比特序列长度必须是8的倍数")
        
        encoded_bytes = bytearray()
        for i in range(0, bit_length, 8):
            byte_bits = mark_bits[i:i+8]
            encoded_bytes.append(int(byte_bits, 2))
        
        token = token_encoder.decode_token(bytes(encoded_bytes), privacy_level)
        return token
    except Exception as e:
        raise ValueError(f"解码标识失败: {str(e)}")

def mark_generate(original_data, token):
    """生成带有标识的数据和标识比特（返回两个值以匹配调用）"""
    mark_bits, original_data, privacy_level = generate(original_data, token)
    # 如果需要嵌入水印，在这里调用Position_mark
    # marked_data = Position_mark.mark_embedding(mark_bits, original_data)


    # 这里输出一下mark_bits长度
    #print(f"标识比特长度: {len(mark_bits)} bits")
    
    return original_data, mark_bits  # 只返回两个值以匹配调用

def mark_identity(mark_bits, original_data):
    """从标识比特识别token，使用原始数据自动计算隐私等级"""
    return decode(mark_bits, original_data)

# 测试函数
def test_decode_with_original_data():
    """测试使用原始数据自动计算隐私等级的解码功能"""
    test_cases = [
        "这是一些示例原始数据" * 10,
        '{"name": "张三", "age": 30, "email": "zhangsan@example.com"}',
        '{"users": [{"name": "张三", "password": "123456"}]}',
        {"username": "testuser", "secret_token": "abc123def456ghi789jk"},
        [{"item": "book", "price": 29.99}, {"item": "pen", "price": 1.99}]
    ]
    
    original_token = "abc123def456ghi789jk"
    
    for i, original_data in enumerate(test_cases):
        print(f"\n=== 测试用例 {i+1} ===")
        print(f"数据类型: {type(original_data).__name__}")
        
        # 生成标识
        data_with_mark, mark_bits = mark_generate(original_data, original_token)
        privacy_level = calculate_privacy_level(original_data)
        print(f"隐私等级: {privacy_level}")
        print(f"标识比特长度: {len(mark_bits)} bits")
        
        # 识别标识（使用原始数据自动计算隐私等级）
        recovered_token = mark_identity(mark_bits, original_data)
        print(f"原始Token: {original_token}")
        print(f"恢复的Token: {recovered_token}")
        print(f"Token匹配: {original_token == recovered_token}")

if __name__ == '__main__':
    test_decode_with_original_data()